import numpy as np
import jupedsim as jps
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

# 参数设置（根据论文数据）
V0 = 1.2  # 初始速度 (m/s)
a = 0.71
b = 0.51
c_down = 1.0
c_up = 2.7
dt = 0.5  # 时间步长 (s)
floor_height = 3.0  # 每层楼高 (m)
num_floors = 6
people_per_floor = 20

# 场景参数
floor_plan_width = 20.0  # 每层楼宽度 (m)
floor_plan_depth = 15.0  # 每层楼深度 (m)
stair_positions = [(5, 5), (15, 10)]  # 楼梯在每层的平面坐标
exit_positions = [(2, 2), (18, 13)]   # 一层出口坐标

class EvacuationSimulator:
    def __init__(self):
        # 创建几何结构
        geometry = []
        for floor in range(num_floors):
            floor_z = floor * floor_height
            # 每层楼的地板区域（矩形）
            floor_poly = [
                (-floor_plan_width/2, -floor_plan_depth/2, floor_z),
                (floor_plan_width/2, -floor_plan_depth/2, floor_z),
                (floor_plan_width/2, floor_plan_depth/2, floor_z),
                (-floor_plan_width/2, floor_plan_depth/2, floor_z)
            ]
            geometry.append(floor_poly)
        
        # 添加楼梯（作为垂直通道）
        for i in range(num_floors - 1):
            stair_z_start = i * floor_height
            stair_z_end = (i + 1) * floor_height
            stair_center = stair_positions[0][:2]  # 假设所有楼梯在同一位置
            stair_poly = [
                (stair_center[0] - 0.5, stair_center[1] - 0.5, stair_z_start),
                (stair_center[0] + 0.5, stair_center[1] - 0.5, stair_z_start),
                (stair_center[0] + 0.5, stair_center[1] + 0.5, stair_z_start),
                (stair_center[0] - 0.5, stair_center[1] + 0.5, stair_z_start),
                (stair_center[0] - 0.5, stair_center[1] - 0.5, stair_z_end),
                (stair_center[0] + 0.5, stair_center[1] - 0.5, stair_z_end),
                (stair_center[0] + 0.5, stair_center[1] + 0.5, stair_z_end),
                (stair_center[0] - 0.5, stair_center[1] + 0.5, stair_z_end)
            ]
            geometry.append(stair_poly)
        
        # 添加出口（在一层）
        for exit_pos in exit_positions:
            exit_poly = [
                (exit_pos[0] - 0.5, exit_pos[1] - 0.5, 0.0),
                (exit_pos[0] + 0.5, exit_pos[1] - 0.5, 0.0),
                (exit_pos[0] + 0.5, exit_pos[1] + 0.5, 0.0),
                (exit_pos[0] - 0.5, exit_pos[1] + 0.5, 0.0)
            ]
            geometry.append(exit_poly)

        # 初始化模拟器（使用默认模型，如SocialForceModel）
        self.sim = jps.Simulation(
            geometry=geometry,
            model=jps.VelocityModel()  # 如果VelocityModel不存在，尝试其他模型
        )

        # 如果VelocityModel不存在，改用SocialForceModel（需检查是否支持）
        # self.sim = jps.Simulation(
        #     geometry=geometry,
        #     model=jps.SocialForceModel()
        # )

        self.agents = []
        self.accumulated_distance = {}
        self.accumulated_time = {}
        
        # 初始化行人
        for floor in range(1, num_floors + 1):
            floor_z = floor * floor_height
            for _ in range(people_per_floor):
                # 随机生成初始位置（避开楼梯区域）
                while True:
                    x = np.random.uniform(-floor_plan_width/2, floor_plan_width/2)
                    y = np.random.uniform(-floor_plan_depth/2, floor_plan_depth/2)
                    if (x, y) not in stair_positions:
                        break
                start_pos = (x, y, floor_z)
                agent = self.sim.add_agent(
                    position=start_pos,
                    radius=0.3,
                    max_speed=V0
                )
                self.accumulated_distance[agent] = 0.0
                self.accumulated_time[agent] = 0.0
                self.agents.append(agent)
    
    def step(self):
        # 更新所有行人状态
        for agent in self.agents:
            current_pos = self.sim.get_position(agent)
            prev_pos = self.sim.get_previous_position(agent)
            if prev_pos is not None:
                distance_moved = np.linalg.norm(np.array(current_pos) - np.array(prev_pos))
                self.accumulated_distance[agent] += distance_moved
                self.accumulated_time[agent] += dt
            
            # 计算当前速度（考虑体力消耗）
            c = c_down if current_pos[2] < prev_pos[2] else c_up
            speed_factor = (1 - a * self.accumulated_distance[agent] - 
                           b * self.accumulated_time[agent]) * np.exp(-0.01 * c)
            current_speed = V0 * max(speed_factor, 0.1)  # 速度不低于0.1m/s
            
            # 更新行人速度
            self.sim.set_agent_max_speed(agent, current_speed)
        
        self.sim.step(dt)
    
    def run(self, max_time=600):
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        ax.set_xlim(-floor_plan_width/2, floor_plan_width/2)
        ax.set_ylim(-floor_plan_depth/2, floor_plan_depth/2)
        ax.set_zlim(0, num_floors * floor_height)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z (Floor)')
        
        lines = [ax.plot([], [], [], 'ro', markersize=5)[0] for _ in self.agents]
        
        def update(frame):
            self.step()
            for i, agent in enumerate(self.agents):
                pos = self.sim.get_position(agent)
                lines[i].set_data([pos[0]], [pos[1]])
                lines[i].set_3d_properties([pos[2]])
            return lines
        
        anim = FuncAnimation(fig, update, frames=int(max_time/dt),
                            interval=int(dt * 1000), blit=False)
        plt.show()

if __name__ == "__main__":
    sim = EvacuationSimulator()
    sim.run()